import time
from gate_api import SpotApi, FuturesApi, Configuration, ApiClient, Order
from datetime import datetime
import logging
import json
import threading
import websocket
import math
#gat 内部套利

class GateArbitrage:
    def __init__(self, api_key: str, api_secret: str):
        # 初始化API客户端
        self.configuration = Configuration(
            key=api_key,
            secret=api_secret,
            host="https://api.gateio.ws/api/v4"
        )
        self.client = ApiClient(self.configuration)
        self.futures_api = FuturesApi(self.client)
        self.spot_api = SpotApi(self.client)
        
        # 价格数据存储
        self.futures_price = {}
        self.spot_price = {}
        
        # websocket连接
        self.futures_ws = None
        self.spot_ws = None
        
        # 初始化日志
        self.logger = logging.getLogger("gate_arbitrage")
        self.logger.setLevel(logging.INFO)
        
        # 只添加文件处理器，不添加控制台处理器
        file_handler = logging.FileHandler(f'gate_arbitrage_{datetime.now().strftime("%Y%m%d")}.log')
        file_handler.setFormatter(logging.Formatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s'
        ))
        self.logger.addHandler(file_handler)
        
        # 添加账户余额存储
        self.spot_balances = {}
        self.futures_balances = {}
        
        # 添加签名相关
        self.api_key = api_key
        self.api_secret = api_secret
        
        # 添加心跳和重连相关的变量
        self.futures_last_ping = 0
        self.spot_last_ping = 0
        self.ping_interval = 5  # 5秒发送一次心跳
        self.reconnect_wait = 5  # 重连等待时间
        self.max_reconnect_attempts = 5  # 最大重连次数
        
        # 添加数据有效性验证的时间戳
        self.futures_balance_timestamp = 0
        self.spot_balance_timestamp = 0
        self.balance_validity_period = 30  # 余额数据有效期(秒)
        
        # 添加错误统计
        self.error_counts = {
            'websocket': 0,
            'api': 0,
            'order': 0,
            'balance': 0
        }
        self.max_error_counts = {
            'websocket': 10,  # websocket错误最大次数
            'api': 5,         # API错误最大次数
            'order': 3,       # 下单错误最大次数
            'balance': 5      # 余额检查错误最大次数
        }
        
        # 添加交易记录管理
        self.trade_records = TradeRecord()
        
        # 添加价格检查参数
        self.price_check_params = {
            'max_price_deviation': 5.0,  # 最大价格偏差(%)
            'min_volume': 1000,          # 最小交易量(USDT)
            'max_slippage': 0.1          # 最大滑点(%)
        }
        
        # 添加仓位管理相关参数
        self.position_limits = {
            'max_total_position': 100,  # 所有仓位总额上限(USDT)
            'max_single_position': 30,   # 单个交易对仓位上限(USDT)
            'min_single_position': 15     # 单个交易对最小仓位(USDT)
        }
        
        # 当前持仓记录
        self.current_positions = {}  # 格式: {symbol: {'amount': float, 'timestamp': float}}

    def generate_sign(self, channel, event, timestamp):
        """生成websocket认证签名"""
        message = 'channel=' + channel + '&' + 'event=' + event + '&' + 'time=' + str(timestamp)
        import hmac
        import hashlib
        h = hmac.new(self.api_secret.encode('utf8'), message.encode('utf8'), hashlib.sha512)
        return h.hexdigest()

    def on_futures_message(self, ws, message):
        """处理合约websocket消息"""
        try:
            data = json.loads(message)
            # self.logger.debug(f"合约消息: {data}")  # 添加调试日志
            
            # 检查是否为更新消息
            if 'event' in data and data['event'] == 'update':
                if data['channel'] == 'futures.tickers':
                    tickers = data['result']  # 这是一个列表
                    for ticker in tickers:  # 遍历列表处理每个ticker
                        symbol = ticker['contract']
                        self.futures_price[symbol] = {
                            'price': float(ticker['last']),
                            'funding_rate': float(ticker['funding_rate']),
                            'timestamp': time.time()
                        }
                        # self.logger.info(f"更新合约价格: {symbol} = {self.futures_price[symbol]}")
            elif 'event' in data and data['event'] == 'subscribe':
                pass
                # self.logger.debug(f"价格消息订阅成功: {data['payload']}")  # 添加调试日志
            elif data['channel'] == 'futures.balances':
                balance = data['result']
                self.futures_balances = {
                    'total': float(balance['total']),
                    'available': float(balance['available'])
                }
        except Exception as e:
            self.logger.error(f"处理合约websocket消息错误: {str(e)}, message: {message}")

    def on_spot_message(self, ws, message):
        """处理现货websocket消息"""
        try:
            data = json.loads(message)
            # self.logger.debug(f"现货消息: {data}")  # 添加调试日志
            
            # 检查是否为更新消息
            if 'event' in data and data['event'] == 'update':
                if data['channel'] == 'spot.tickers':
                    ticker = data['result']
                    symbol = ticker['currency_pair']
                    self.spot_price[symbol] = {
                        'price': float(ticker['last']),
                        'timestamp': time.time()
                    }
                    # self.logger.info(f"更新现货价格: {symbol} = {self.spot_price[symbol]}")
            elif 'event' in data and data['event'] == 'subscribe':
                pass
            elif data['channel'] == 'spot.balances':
                for balance in data['result']:
                    currency = balance['currency']
                    self.spot_balances[currency] = {
                        'available': float(balance['available']),
                        'locked': float(balance['locked'])
                    }
        except Exception as e:
            self.logger.error(f"处理现货websocket消息错误: {str(e)}, message: {message}")

    def on_futures_error(self, ws, error):
        """处理合约websocket错误"""
        self.logger.error(f"合约Websocket错误: {str(error)}")

    def on_futures_close(self, ws, close_status_code, close_msg):
        """处理合约websocket关闭"""
        self.logger.warning("合约Websocket连接关闭")
        self.reconnect_futures()

    def on_futures_open(self, ws):
        """处理合约websocket连接打开"""
        self.logger.info("合约Websocket连接成功")
        self.authenticate_futures()

    def on_spot_error(self, ws, error):
        """处理现货websocket错误"""
        self.logger.error(f"现货Websocket错误: {str(error)}")

    def on_spot_close(self, ws, close_status_code, close_msg):
        """处理现货websocket关闭"""
        self.logger.warning("现货Websocket连接关闭")
        self.reconnect_spot()

    def on_spot_open(self, ws):
        """处理现货websocket连接打开"""
        self.logger.info("现货Websocket连接成功")
        self.authenticate_spot()

    def reconnect_futures(self):
        """重连合约websocket"""
        attempts = 0
        while attempts < self.max_reconnect_attempts:
            try:
                self.logger.info(f"尝试重连合约Websocket, 第{attempts + 1}次")
                time.sleep(self.reconnect_wait)
                self.start_futures_ws()
                return
            except Exception as e:
                attempts += 1
                self.logger.error(f"合约Websocket重连失败: {str(e)}")
        self.logger.error("合约Websocket重连失败次数超过最大限制")

    def reconnect_spot(self):
        """重连现货websocket"""
        attempts = 0
        while attempts < self.max_reconnect_attempts:
            try:
                self.logger.info(f"尝试重连现货Websocket, 第{attempts + 1}次")
                time.sleep(self.reconnect_wait)
                self.start_spot_ws()
                return
            except Exception as e:
                attempts += 1
                self.logger.error(f"现货Websocket重连失败: {str(e)}")
        self.logger.error("现货Websocket重连失败次数超过最大限制")

    def send_ping(self):
        """发送心跳"""
        current_time = time.time()
        
        # 合约心跳
        if current_time - self.futures_last_ping >= self.ping_interval:
            try:
                self.futures_ws.send(json.dumps({"time": int(current_time), "channel": "futures.ping"}))
                self.futures_last_ping = current_time
            except Exception as e:
                self.logger.error(f"发送合约心跳失败: {str(e)}")
                self.reconnect_futures()
        
        # 现货心跳
        if current_time - self.spot_last_ping >= self.ping_interval:
            try:
                self.spot_ws.send(json.dumps({"time": int(current_time), "channel": "spot.ping"}))
                self.spot_last_ping = current_time
            except Exception as e:
                self.logger.error(f"发送现货心跳失败: {str(e)}")
                self.reconnect_spot()

    def is_balance_valid(self):
        """检查余额数据是否有效"""
        current_time = time.time()
        futures_valid = (current_time - self.futures_balance_timestamp) <= self.balance_validity_period
        spot_valid = (current_time - self.spot_balance_timestamp) <= self.balance_validity_period
        return futures_valid and spot_valid

    def start_websockets(self, futures_symbol: str, spot_symbol: str):
        """启动websocket连接"""
        def futures_ws_thread():
            while True:
                try:
                    # websocket.enableTrace(True)  # 注释掉或删除这行
                    self.futures_ws = websocket.WebSocketApp(
                        "wss://fx-ws.gateio.ws/v4/ws/usdt",
                        on_message=self.on_futures_message,
                        on_error=self.on_futures_error,
                        on_close=self.on_futures_close,
                        on_open=self.on_futures_open
                    )
                    self.futures_ws.run_forever()
                except Exception as e:
                    self.logger.error(f"合约Websocket运行错误: {str(e)}")
                    time.sleep(self.reconnect_wait)

        def spot_ws_thread():
            while True:
                try:
                    # websocket.enableTrace(True)  # 注释掉或删除这行
                    self.spot_ws = websocket.WebSocketApp(
                        "wss://api.gateio.ws/ws/v4/",
                        on_message=self.on_spot_message,
                        on_error=self.on_spot_error,
                        on_close=self.on_spot_close,
                        on_open=self.on_spot_open
                    )
                    self.spot_ws.run_forever()
                except Exception as e:
                    self.logger.error(f"现货Websocket运行错误: {str(e)}")
                    time.sleep(self.reconnect_wait)

        # 启动websocket线程
        threading.Thread(target=futures_ws_thread, daemon=True).start()
        threading.Thread(target=spot_ws_thread, daemon=True).start()
        
        # 等待连接建立
        time.sleep(3)

        # 合约认证和订阅
        timestamp = int(time.time())
        sign = self.generate_sign('futures.balances', 'subscribe', timestamp)
        
        # 合约认证
        self.futures_ws.send(json.dumps({
            "time": timestamp,
            "channel": "futures.auth",
            "event": "auth",
            "payload": [self.api_key, sign]
        }))

        # 订阅合约行情和余额
        self.futures_ws.send(json.dumps({
            "time": timestamp,
            "channel": "futures.tickers",
            "event": "subscribe",
            "payload": [futures_symbol]
        }))
        
        # 修改合约余额订阅方式
        self.futures_ws.send(json.dumps({
            "time": timestamp,
            "channel": "futures.usertrades",  # 改为用户交易频道
            "event": "subscribe",
            "payload": [""]  # 空字符串表示订阅所有合约
        }))

        # 请求合约账户余额
        self.futures_ws.send(json.dumps({
            "time": timestamp,
            "channel": "futures.query_accounts",  # 查询账户余额
            "event": "query",
            "payload": []
        }))

        # 现货认证和订阅
        timestamp = int(time.time())
        sign = self.generate_sign('spot.balances', 'subscribe', timestamp)
        
        # 现货认证
        self.spot_ws.send(json.dumps({
            "time": timestamp,
            "channel": "spot.auth",
            "event": "auth",
            "payload": [self.api_key, sign]
        }))

        # 订阅现货行情和余额
        self.spot_ws.send(json.dumps({
            "time": timestamp,
            "channel": "spot.tickers",
            "event": "subscribe",
            "payload": [spot_symbol]
        }))
        
        self.spot_ws.send(json.dumps({
            "time": timestamp,
            "channel": "spot.balances",
            "event": "subscribe",
            "payload": []
        }))

    def handle_error(self, error_type: str, error: Exception, critical: bool = False):
        """统一错误处理"""
        self.error_counts[error_type] += 1
        error_msg = f"{error_type.upper()} Error: {str(error)}"
        
        if critical:
            self.logger.error(error_msg, exc_info=True)
        else:
            self.logger.warning(error_msg)
            
        # 检查错误次数是否超过限制
        if self.error_counts[error_type] >= self.max_error_counts[error_type]:
            self.logger.critical(f"{error_type} 错误次数超过限制，程序将退出")
            raise SystemExit(f"{error_type} 错误次数超过限制")

    def get_price_difference(self, futures_symbol: str, spot_symbol: str):
        """获取永续合约和现货的价格差异"""
        try:
            if futures_symbol not in self.futures_price or spot_symbol not in self.spot_price:
                self.logger.warning(f"价格数据不完整: futures={futures_symbol in self.futures_price}, "
                                  f"spot={spot_symbol in self.spot_price}")
                return None
                
            futures_data = self.futures_price[futures_symbol]
            spot_price = self.spot_price[spot_symbol]
            
            # 价格有效性检查
            if futures_data['price'] <= 0 or spot_price['price'] <= 0:
                self.logger.warning(f"检测到无效价格: futures={futures_data['price']}, spot={spot_price}")
                return None
            
            # 计算价格差异百分比
            price_diff = (futures_data['price'] - spot_price['price']) / spot_price['price']
            
            result = {
                'futures_price': futures_data['price'],
                'spot_price': spot_price['price'],
                'price_diff': price_diff,
                'funding_rate': futures_data['funding_rate'],
                'futures_symbol':futures_symbol,
                'spot_symbol':spot_symbol
            }
            
            # self.logger.debug(f"价格差异计算结果: {result}")
            return result
            
        except Exception as e:
            self.handle_error('api', e)
            return None

    def check_balance(self, currency: str, required_amount: float):
        """检查账户余额是否足够"""
        try:
            # 获取现货余额
            spot_accounts = self.spot_api.list_spot_accounts(currency=currency)
            spot_available = float(spot_accounts[0].available) if spot_accounts else 0
            
            # 获取合约余额 - 修正方法名
            futures_accounts = self.futures_api.list_futures_accounts(settle='usdt')
            futures_available = float(futures_accounts.total)  # 使用total字段获取可用余额
            
            self.logger.info(f"现货可用余额: {spot_available} {currency}")
            self.logger.info(f"合约可用余额: {futures_available} USDT")
            
            # 更新本地缓存
            self.spot_balances[currency] = {
                'available': spot_available,
                'timestamp': time.time()
            }
            
            self.futures_balances = {
                'available': futures_available,
                'timestamp': time.time()
            }
            
            return spot_available >= required_amount and futures_available >= required_amount
            
        except Exception as e:
            self.logger.error(f"检查余额失败: {str(e)}")
            return False

    def verify_price(self, price_info: dict) -> bool:
        """验证价格有效性"""
        try:
            futures_symbol = price_info['futures_symbol']
            spot_symbol = price_info['spot_symbol']
            futures_price = price_info['futures_price']
            spot_price = price_info['spot_price']
            
            # 获取原始数据中的时间戳
            futures_timestamp = self.futures_price[futures_symbol]['timestamp']
            spot_timestamp = self.spot_price[spot_symbol]['timestamp']
            current_time = time.time()
            
            # 检查数据时效性 (默认5秒内的数据有效)
            max_data_age = 5  # 数据有效期（秒）
            if current_time - futures_timestamp > max_data_age:
                self.logger.warning(f"期货价格数据过期: {futures_symbol}, 数据时间: {futures_timestamp}, 当前时间: {current_time}")
                return False
            
            if current_time - spot_timestamp > max_data_age:
                self.logger.warning(f"现货价格数据过期: {spot_symbol}, 数据时间: {spot_timestamp}, 当前时间: {current_time}")
                return False
            
            # 检查价格偏差
            price_deviation = abs((futures_price - spot_price) / spot_price)
            if price_deviation > self.price_check_params['max_price_deviation']:
                self.logger.warning(f"价格偏差过大: {price_deviation}%")
                return False
            
            # 获取深度数据进行验证
            futures_depth = self.futures_price.get(futures_symbol, {})
            spot_depth = self.spot_price.get(spot_symbol, {})
            
            # 验证交易量
            if not self.verify_volume(futures_depth, spot_depth):
                return False
            
            # 验证滑点
            if not self.verify_slippage(futures_depth, spot_depth, futures_price, spot_price):
                return False
            
            return True
            
        except Exception as e:
            self.logger.error(f"价格验证失败: {str(e)}")
            return False

    def verify_volume(self, futures_depth: dict, spot_depth: dict) -> bool:
        """验证交易量"""
        try:
            # 这里添加具体的交易量验证逻辑
            return True
        except Exception as e:
            self.logger.error(f"交易量验证失败: {str(e)}")
            return False
            
    def verify_slippage(self, futures_depth: dict, spot_depth: dict, 
                       futures_price: float, spot_price: float) -> bool:
        """验证滑点"""
        try:
            # 这里添加具体的滑点验证逻辑
            return True
        except Exception as e:
            self.logger.error(f"滑点验证失败: {str(e)}")
            return False

    def place_arbitrage_orders(self, futures_symbol: str, spot_symbol: str, 
                             amount: float, is_long_futures: bool = True):
        """下单"""
        try:

            # 检查仓位限制
            if not self.check_position_limits(futures_symbol, amount):
                raise Exception("不满足仓位限制条件")

            price_info = self.get_price_difference(futures_symbol, spot_symbol)
            
            # 根据价差确定交易方向
            price_diff = price_info['price_diff']  # (futures_price - spot_price) / spot_price
            
            # 当合约价格高于现货价格时(price_diff > 0)，做空合约
            # 当合约价格低于现货价格时(price_diff < 0)，做多合约
            is_long_futures = price_diff < 0
            
            self.trade_records.start_trade(
                futures_symbol, spot_symbol, 
                price_info,
                is_long_futures
            )
            
            # 获取合约信息
            contract_info = self.futures_api.get_futures_contract(settle='usdt', contract=futures_symbol)
            current_price = self.futures_price[futures_symbol]['price']
            
            # 计算合约张数
            contract_size = float(contract_info.quanto_multiplier) if hasattr(contract_info, 'quanto_multiplier') else 1
            value_per_contract = contract_size * current_price
            contract_amount = math.ceil(amount /current_price /contract_size)
            
            # 根据方向设置size的正负
            if not is_long_futures:  # 做空时size为负数
                contract_amount = -contract_amount
            
            actual_value = abs(contract_amount) * value_per_contract
            
            self.logger.info(f"下单信息: 目标USDT金额={amount}, 实际USDT金额={actual_value}, "
                            f"合约价格={current_price}, 每张合约大小={contract_size}, "
                            f"每张合约价值={value_per_contract}, 合约张数={contract_amount}")
            
            # 检查偏差
            max_deviation = 0.1
            if abs(actual_value - amount) / amount > max_deviation:
                raise Exception(f"实际下单金额 {actual_value} USDT 与目标金额 {amount} USDT 偏差过大")
            
            # 合约下单 - 使用size的正负来控制方向
            futures_order_request = {
                'contract': futures_symbol,
                'size': contract_amount,  # 正数做多，负数做空
                'price': '0',  # 市价单
                'tif': 'ioc',
                'type': 'market',
                'close': False,
                'reduce_only': False
            }
            
            futures_order = self.futures_api.create_futures_order(
                settle='usdt',
                futures_order=futures_order_request
            )
            
            self.logger.info(f"合约下单成功: {futures_order}")
            
            # 现货下单 - 与合约方向相反
            spot_side = 'sell' if is_long_futures else 'buy'  # 合约做多时现货卖出
            spot_order_request = {
                'currency_pair': spot_symbol,
                'amount': str(abs(actual_value)),  # 使用合约张数 * 每张合约的币数量
                'side': spot_side,
                'type': 'market',
                'time_in_force': 'ioc'  # 改为 IOC (Immediate or Cancel)
            }

            spot_order = self.spot_api.create_order(
                order=spot_order_request  # 传入订单参数
            )

            self.logger.info(f"现货下单成功: {spot_order}")
            
            # 验证订单状态
            if not self.verify_orders(futures_order, spot_order):
                raise Exception("订单验证失败")
            
            if futures_order and spot_order:
                self.trade_records.update_trade(
                    futures_order=futures_order,
                    spot_order=spot_order,
                    status='opened'
                )
            else:
                self.trade_records.update_trade(
                    status='failed',
                    error='Order placement failed'
                )
            

            
            # 更新持仓记录
            if futures_order and spot_order:
                self.current_positions[futures_symbol] = {
                    'amount': actual_value,
                    'timestamp': time.time()
                }
            
            return futures_order, spot_order
            
        except Exception as e:
            self.trade_records.update_trade(
                status='error',
                error=str(e)
            )
            self.handle_error('order', e, critical=True)
            return None, None

    def verify_orders(self, futures_order, spot_order):
        """验证订单状态"""
        try:
            if not futures_order or not spot_order:
                self.logger.error("订单对象为空")
                return False
                
            # 这里可以添加更多的订单验证逻辑
            # 例如检查成交状态、成交价格等
            
            return True
        except Exception as e:
            self.logger.error(f"订单验证失败: {str(e)}")
            return False

    def run_arbitrage(self, futures_symbol: str, spot_symbol: str, min_diff: float, 
                     amount: float, check_interval: float = 1.0):
        """运行套利程序"""
        try:
            # 启动websocket
            self.start_websockets(futures_symbol, spot_symbol)
            
            while True:
                try:
                    # 发送心跳
                    self.send_ping()
                    
                    # 检查价格差异
                    price_info = self.get_price_difference(futures_symbol, spot_symbol)
                    
                    if price_info:
                        # 验证价格有效性
                        if not self.verify_price(price_info):
                            continue
                            
                        self.logger.info(f"市场状态: 价格差异={price_info['price_diff']:.2f}%, "
                                       f"资金费率={price_info['funding_rate']:.4f}%, "
                                       f"合约价格={price_info['futures_price']}, "
                                       f"现货价格={price_info['spot_price']}")
                        
                        # 如果价格差异超过阈值，执行套利
                        if abs(price_info['price_diff']) >= min_diff:
                            # 二次确认价格
                            confirm_price_info = self.get_price_difference(futures_symbol, spot_symbol)
                            if not confirm_price_info or abs(confirm_price_info['price_diff'] - price_info['price_diff']) > 0.1:
                                self.logger.warning("价格二次确认失败")
                                continue
                                
                            self.logger.info(f"发现套利机会: 价差={price_info['price_diff']:.2f}%")
                            
                            # 检查余额
                            if self.check_balance('USDT', amount * price_info['spot_price']):
                                # 确定交易方向
                                is_long_futures = price_info['price_diff'] < 0
                                
                                # 执行下单
                                futures_order, spot_order = self.place_arbitrage_orders(
                                    futures_symbol, spot_symbol, amount, is_long_futures
                                )
                                
                                if futures_order and spot_order:
                                    self.logger.info(f"套利执行成功! "
                                                   f"合约订单ID: {futures_order.id}, "
                                                   f"现货订单ID: {spot_order.id}")
                                    # 重置错误计数
                                    self.error_counts['order'] = 0
                                
                                # 保存交易记录
                                self.trade_records.save_to_db(self.mysql_utils)
                
                    time.sleep(check_interval)
                    
                except Exception as e:
                    self.handle_error('api', e)
                    time.sleep(check_interval)
                    
        except Exception as e:
            self.logger.critical("套利程序发生严重错误", exc_info=True)
            raise

    def get_current_positions(self, symbol: str):
        """获取当前持仓情况"""
        try:
            # 获取合约持仓
            futures_positions = self.futures_api.list_positions(settle='usdt')
            total_position = 0
            
            for position in futures_positions:
                if position.contract == symbol:
                    size = float(position.size)
                    entry_price = float(position.entry_price)
                    position_value = abs(size) * entry_price
                    total_position += position_value
                    
                    self.current_positions[symbol] = {
                        'amount': position_value,
                        'timestamp': time.time()
                    }
                    
            return total_position
            
        except Exception as e:
            self.logger.error(f"获取持仓信息失败: {str(e)}")
            return None

    def check_position_limits(self, symbol: str, new_position_size: float) -> bool:
        """检查是否满足仓位限制"""
        try:
            # 获取当前持仓
            current_position = self.get_current_positions(symbol) or 0
            
            # 计算所有交易对的总持仓
            total_position = sum(pos['amount'] for pos in self.current_positions.values())
            
            # 检查是否超过总仓位限制
            if total_position + new_position_size > self.position_limits['max_total_position']:
                self.logger.warning(f"总仓位 {total_position + new_position_size} USDT 超过限制 "
                                  f"{self.position_limits['max_total_position']} USDT")
                return False
            
            # 检查单个交易对仓位限制
            if current_position + new_position_size > self.position_limits['max_single_position']:
                self.logger.warning(f"单个交易对仓位 {current_position + new_position_size} USDT 超过限制 "
                                  f"{self.position_limits['max_single_position']} USDT")
                return False
            
            # 检查最小仓位要求
            if new_position_size < self.position_limits['min_single_position']:
                self.logger.warning(f"开仓金额 {new_position_size} USDT 小于最小要求 "
                                  f"{self.position_limits['min_single_position']} USDT")
                return False
            
            return True
            
        except Exception as e:
            self.logger.error(f"检查仓位限制失败: {str(e)}")
            return False

class TradeRecord:
    def __init__(self):
        self.trades = []
        self.current_trade = None
        
    def start_trade(self, futures_symbol: str, spot_symbol: str, price_info: dict, is_long_futures: bool):
        """开始一笔新交易"""
        self.current_trade = {
            'start_time': datetime.now(),
            'futures_symbol': futures_symbol,
            'spot_symbol': spot_symbol,
            'futures_price': price_info['futures_price'],
            'spot_price': price_info['spot_price'],
            'price_diff': price_info['price_diff'],
            'funding_rate': price_info['funding_rate'],
            'direction': 'long_futures' if is_long_futures else 'short_futures',
            'status': 'opening',
            'orders': {
                'futures': None,
                'spot': None
            },
            'profit': None,
            'error': None
        }
        
    def update_trade(self, futures_order=None, spot_order=None, status=None, error=None):
        """更新当前交易信息"""
        if not self.current_trade:
            return
            
        if futures_order:
            self.current_trade['orders']['futures'] = futures_order
        if spot_order:
            self.current_trade['orders']['spot'] = spot_order
        if status:
            self.current_trade['status'] = status
        if error:
            self.current_trade['error'] = error
            
    def complete_trade(self, profit=None):
        """完成当前交易"""
        if self.current_trade:
            self.current_trade['end_time'] = datetime.now()
            self.current_trade['profit'] = profit
            self.trades.append(self.current_trade)
            self.current_trade = None
            
    def save_to_db(self, mysql_utils):
        """保存交易记录到数据库"""
        if not self.current_trade:
            return
            
        sql = """
        INSERT INTO trade_records (
            start_time, end_time, futures_symbol, spot_symbol,
            futures_price, spot_price, price_diff, funding_rate,
            direction, status, futures_order_id, spot_order_id,
            profit, error
        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
        """
        
        params = (
            self.current_trade['start_time'],
            self.current_trade.get('end_time'),
            self.current_trade['futures_symbol'],
            self.current_trade['spot_symbol'],
            self.current_trade['futures_price'],
            self.current_trade['spot_price'],
            self.current_trade['price_diff'],
            self.current_trade['funding_rate'],
            self.current_trade['direction'],
            self.current_trade['status'],
            self.current_trade['orders']['futures']['id'] if self.current_trade['orders']['futures'] else None,
            self.current_trade['orders']['spot']['id'] if self.current_trade['orders']['spot'] else None,
            self.current_trade['profit'],
            self.current_trade['error']
        )
        
        try:
            mysql_utils.execute(sql, params)
        except Exception as e:
            self.logger.error(f"保存交易记录失败: {str(e)}")

def main():
    # API配置
    api_key = "3e8bf1c102bedfdea327212db1b466be"
    api_secret = "d8d6024fa11a053f71d7dd20c1ce0713335d0b4e88e80ba11a844c00684e1d84"
    
    # 创建套利实例
    arbitrage = GateArbitrage(api_key, api_secret)
    
    # 运行套利程序
    arbitrage.run_arbitrage(
        futures_symbol="WOLF_USDT",  # 永续合约符号
        spot_symbol="WOLF_USDT",     # 现货交易对
        min_diff=0.001,               # 最小价差阈值 (0.5%)
        amount=15,                # 交易数量
        check_interval=1.0          # 检查间隔(秒)
    )

if __name__ == "__main__":
    main()
